from mxnet import ndarray as nd

print(nd.zeros((2,3)))
print(nd.array([[1,2]]))

print(nd.random_normal(0, 1, shape=(3, 3)))

import mxnet.autograd as ag

# f = 2 * x^2
x = nd.array([[1, 2], [3, 4]])
x.attach_grad()

with ag.record():
    y = x * 2
    z = y * x

z.backward()
print(x.grad)


# grad for control flow
def f(a):
    b = a * 2
    while nd.norm(b).asscalar() < 1000:
        b = b * 2

    if nd.sum(b).asscalar() > 0:
        c = b
    else:
        c = 100 * b
    return c


a = nd.random_normal(shape=3)
a.attach_grad()
with ag.record():
    c = f(a)
c.backward()

print(a.grad)
print(a.grad == c/a)

# head gradient
with ag.record():
    y = x * 2
    z = y * x

head_gradient = nd.array([[10, 1.], [.1, .01]])
z.backward(head_gradient)
print(x.grad)
